
from typing import Any, Callable, Dict, List, Optional
from dataclasses import dataclass
import functools
import asyncio
import inspect

##################################################################
# 公用的的函数： 比如，以保证中文输出的帮助界面能够对齐的函数等
##################################################################
def get_visual_width(s: str) -> int:
    """
    计算字符串的视觉宽度，中文字符计为2，其他字符计为1。
    """
    width = 0
    for char in s:
        # \u4e00-\u9fa5 是中文字符的主要Unicode范围
        # \uff00-\uffef 是全角字符（包括全角标点）的范围
        if '\u4e00' <= char <= '\u9fa5' or '\uff00' <= char <= '\uffef':
            width += 2
        else:
            width += 1
    return width


def create_aligned_intro(lines, total_width=80) -> str:
    """
    创建右对齐的欢迎语字符串。
    """
    total_width = max(len(x) for x in lines) * 2 - 2

    # 1. 构建字符串
    border_top = "╔" + "═" * total_width + "╗"
    border_bottom = "╚" + "═" * total_width + "╝"
    empty_line = "║" + " " * total_width + "║"

    intro_parts = [
        "\n",
        border_top,
        empty_line
    ]

    # 2. 动态计算每行的填充并添加到列表中
    for text in lines:
        visual_width = get_visual_width(text)
        padding = total_width - visual_width
        # 确保填充不为负数
        padding = max(0, padding)

        line_str = f"║{text}{' ' * padding}║"
        intro_parts.append(line_str)

    intro_parts.append(empty_line)
    intro_parts.append(border_bottom)
    intro_parts.append("\n")

    return "\n".join(intro_parts)


##################################################################
# Tool System
##################################################################

@dataclass # 简化类定义
class Tool:
    name: str
    description: str
    function: Callable
    required_args: List[str]
    optional_args: Optional[Dict[str, Any]] = None

    def validate_args(self, **kwargs):
        for arg in self.required_args:
            if arg not in kwargs:
                raise ValueError(f"Missing required argument: {arg}")

    async def __call__(self, **kwargs):
        self.validate_args(**kwargs)
        result = self.function(**kwargs)
        if asyncio.iscoroutine(result):
            return await result
        return result


class ToolSystem:
    def __init__(self):
        self._tools: Dict[str, Tool] = {}
        self.lines = []

    def register_tool(self, tool: Tool):
        """注册新工具到系统"""
        if tool.name in self._tools:
            raise ValueError(f"Tool {tool.name} already exists")
        self._tools[tool.name] = tool

    def get_tool(self, name: str) -> Tool:
        """获取指定工具"""
        if name not in self._tools:
            raise ValueError(f"Tool {name} not found")
        return self._tools[name]

    def list_tools(self) -> List[Dict[str, str]]:
        """列出所有可用工具"""
        return [{"name": t.name, "description": t.description}
                for t in self._tools.values()]

    def get_tools(self):
        return self._tools

    def has_tool(self, name: str) -> bool:
        return name in self._tools

    @staticmethod
    def _parse_mcp_tool_schema(tool_info) -> tuple:
        """
        从 fastmcp 的工具信息对象中解析出名称、描述、必需参数和可选参数。
        """
        # 1. 获取基本信息
        name = tool_info.name
        description = tool_info.description or ""  # 使用 .description 属性，并提供默认值

        # 2. 初始化参数列表
        required_args = []
        optional_args_with_defaults = {}

        # 3. 安全地获取 inputSchema 和其内部结构
        input_schema = getattr(tool_info, 'inputSchema', {})
        if not input_schema:  # 如果没有输入 schema，则没有参数
            return name, description, [], {}

        properties = input_schema.get('properties', {})
        required_names = input_schema.get('required', [])

        if not properties:  # 如果 schema 中没有 properties，则没有参数
            return name, description, [], {}

        # 4. 遍历所有定义的参数，并进行分类
        for param_name, param_details in properties.items():
            if param_name in required_names:
                # 这是必需参数
                required_args.append(param_name)
            else:
                # 这是可选参数，检查它是否有默认值
                if 'default' in param_details:
                    optional_args_with_defaults[param_name] = param_details['default']
                # 注意：即使没有默认值，它仍然是可选的，
                # 但我们的 Tool 类目前只关心有默认值的可选参数。

        return name, description, required_args, optional_args_with_defaults

    async def load_tools_from_mcp(self, mcp_url: str):
        """
        从远程 MCP 服务加载工具并注册。
        参数:
            endpoint: MCP 服务地址, 例如 "http://127.0.0.1:9000"
        """
        from fastmcp import Client  # 引用放在这里：考虑与前3次课程代码的兼容性

        if mcp_url is None:
            mcp_url = "http://127.0.0.1:9000"
        print(f"🔗 正在连接 MCP 服务: {mcp_url}")
        try:
            async with Client(f"{mcp_url}/mcp") as client:
                tools_info_list  = await client.list_tools()
                if not tools_info_list:
                    print(f"ℹ️ 在 {mcp_url} 上未发现可用工具。")
                    return
                print(f"🎯 从 MCP 服务发现 {len(tools_info_list)} 个工具，正在注册...")
                for tool_info in tools_info_list:
                    name, desc, required, optional = self._parse_mcp_tool_schema(tool_info)

                    # 创建一个独立的 async 包装器，每次调用时再开连接
                    async def mcp_tool_wrapper(_tool_name=name, **kwargs):
                        async with Client(f"{mcp_url}/mcp") as inner_client:
                            # return await inner_client.call_tool(name, kwargs)
                            return await inner_client.call_tool(_tool_name, kwargs)

                    #functools.update_wrapper(mcp_tool_wrapper, client.call_tool)
                    # 注册到系统
                    try:
                        self.register_tool(
                            Tool(name=name, description=desc, function=mcp_tool_wrapper, required_args=required,
                                 optional_args=optional)
                        )
                        print(f"✅ 已注册远程 MCP 工具: {name}, {desc}, {required}, {optional}")
                    except ValueError as e:
                        print(f"⚠️ 注册失败：{e}，继续下一个。")


        except Exception as e:
            print(f"MCP server {mcp_url} 连接失败， {e}")
            return

        print(f"🎯 从 {mcp_url} 成功加载 {len(tools_info_list)} 个工具。")

    async def run_tool(self, tool_name: str, args=None):
        """
        运行注册的工具。
        ✅ 支持以下参数形式：
            dict  -> {"path": "."}
            str   -> "."
            list  -> [".", "other"]
            None  -> 无参数
        自动匹配 required_args / optional_args。
        """
        if tool_name not in self._tools:
            return f"❌ 工具 '{tool_name}' 未注册。"

        tool = self._tools[tool_name]
        func = tool.function

        # ------------------------
        # 参数类型统一转换
        # ------------------------
        if args is None:
            kwargs = {}
        elif isinstance(args, dict):
            kwargs = args
        elif isinstance(args, (list, tuple)):
            arg_names = tool.required_args + list(tool.optional_args or [])
            kwargs = {k: v for k, v in zip(arg_names, args)}
        elif isinstance(args, str):
            if tool.required_args:
                kwargs = {tool.required_args[0]: args}
            elif tool.optional_args:
                first_key = next(iter(tool.optional_args))
                kwargs = {first_key: args}
            else:
                kwargs = {}
        else:
            return f"⚠️ 不支持的参数类型: {type(args).__name__}"

        # ------------------------
        # 参数验证
        # ------------------------
        missing_args = [a for a in tool.required_args if a not in kwargs]
        if missing_args:
            return f"⚠️ 缺少必需参数: {missing_args}"

        all_allowed = set(tool.required_args) | set(tool.optional_args or {})
        unknown_args = [k for k in kwargs if k not in all_allowed]
        if unknown_args:
            return f"⚠️ 存在未知参数: {unknown_args}"

        # ------------------------
        # 执行函数（同步/异步兼容）
        # ------------------------
        try:
            if inspect.iscoroutinefunction(func):
                result = await func(**kwargs)
            else:
                loop = asyncio.get_event_loop()
                result = await loop.run_in_executor(None, lambda: func(**kwargs))

            return f"✅ 工具 '{tool_name}' 执行成功：{result}"

        except Exception as e:
            return f"❌ 工具 '{tool_name}' 执行失败：{type(e).__name__} - {e}"

##################################################################
# 装饰器：简化工具注册
##################################################################
def tool(tool_system: ToolSystem, name: str, description: str, required_args: List[str], optional_args: Optional[Dict[str, Any]] = None):
    """
    用于注册外部工具函数的装饰器。
    - name: 工具名
    - description: 工具描述
    - required_args: 必须参数列表
    - optional_args: 可选参数字典
    """
    def decorator(func: Callable):
        # 将原函数包装为 Tool 并注册
        new_tool = Tool(
            name=name,
            description=description,
            function=func,
            required_args=required_args,
            optional_args=optional_args
        )
        tool_system.register_tool(new_tool)

        @functools.wraps(func)
        def wrapper(*args, **kwargs):
            return func(*args, **kwargs)

        return wrapper
    return decorator


class UIAgent:
    def __init__(self, registry: ToolSystem):
        self.command_registry = registry

    def show_intro(self):
        try:
            return create_aligned_intro(self.command_registry.lines)
        except Exception:
            return "\n".join(self.command_registry.lines)

    def show_help(self, command=None):
        if command and self.command_registry.has_tool(command):
            t = self.command_registry.get_tool(command)
            usage = f"{t.name} {' '.join(t.required_args)}"
            return f"命令: {t.name}\n功能: {t.description}\n用法: {usage}"
        else:
            return self.show_intro()